package flowdroid.taint;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import soot.SootMethod;
import soot.jimple.infoflow.results.ResultSourceInfo;
import soot.jimple.infoflow.solver.cfg.IInfoflowCFG;

import java.util.*;
import java.util.function.Function;
import java.util.stream.Collectors;

public class ResultFilter implements Function<Set<ResultSourceInfo>, Set<ResultSourceInfo>> {
    // 数据流图
    private final IInfoflowCFG cfg;
    private final Logger logger = LoggerFactory.getLogger(getClass());
    private boolean openReconstruction = true;
    private int failCnt = 0;

    public ResultFilter(IInfoflowCFG cfg) {
        this.cfg = cfg;
    }

    /**
     * 对集合去重
     *
     * @param sources the function argument
     * @return set that drops duplicate
     */
    @Override
    public Set<ResultSourceInfo> apply(Set<ResultSourceInfo> sources) {
        Set<ResultSourceInfo> ret = new HashSet<>();

        for (ResultSourceInfo current : sources) {
            if (ret.isEmpty()) {
                ret.add(current);
                continue;
            }
            boolean isSimilarFlow = false;
            for (ResultSourceInfo visitedSource : ret) {
                if (compareByVisitedMethod(visitedSource, current) == 0) {
                    isSimilarFlow = true;
                    break;
                }
            }
            if (!isSimilarFlow) {
                ret.add(current);
            }
        }
        return ret;
    }

    private int compareByVisitedMethod(ResultSourceInfo source1, ResultSourceInfo source2) {
        Set<SootMethod> visitedMethods1 = tryGetAllFlowMethods(source1);
        Set<SootMethod> visitedMethods2 = tryGetAllFlowMethods(source2);

        if (visitedMethods1.size() < visitedMethods2.size()) {
            return -1;
        } else if (visitedMethods1.size() > visitedMethods2.size()) {
            return 1;
        } else {
            Iterator<SootMethod> iter1 = visitedMethods1.stream().iterator();
            Iterator<SootMethod> iter2 = visitedMethods2.stream().iterator();
            while (iter1.hasNext() && iter2.hasNext()) {
                int i = iter1.next().getSignature().compareTo(iter2.next().getSignature());
                if (i != 0)
                    return i;
            }
        }
        return 0;
    }

    private Set<SootMethod> tryGetAllFlowMethods(ResultSourceInfo resultSourceInfo) {
        Set<SootMethod> flowMethods = new HashSet<>();
        flowMethods.add(cfg.getMethodOf(resultSourceInfo.getStmt()));
        try {
            if (resultSourceInfo.getPath() != null && openReconstruction) {
                LinkedHashSet<SootMethod> tmp = Arrays.stream(resultSourceInfo.getPath())
                        .map(cfg::getMethodOf).sorted(Comparator.comparing(SootMethod::getSignature))
                        .collect(Collectors.toCollection(LinkedHashSet::new));
                flowMethods.addAll(tmp);
            }
        } catch (Exception e) {
            logger.error(e.getMessage());
            failCnt++;
            if (failCnt > 3) {
                openReconstruction = false;
            }
        }
        return flowMethods;
    }
}

